Fix LinearRegressor heap-buffer-overflow from undersized coefficients#27964
Fix LinearRegressor heap-buffer-overflow from undersized coefficients#27964bmehta001 wants to merge 2 commits into
Conversation
There was a problem hiding this comment.
Pull request overview
This PR hardens several operators against invalid shapes/attributes that could lead to out-of-bounds reads (notably LinearRegressor), and adds additional runtime validation in CPU/CUDA/WebGPU and JS implementations.
Changes:
- Add runtime validation for
LinearRegressorthatcoefficientsmatchestargets * num_features, plus a regression test for the mismatch case. - Add additional validation in
Slice/Pad(negative output dims, axis range) andSplit(non-negative split sizes) across multiple execution providers/backends. - Fix JS
normalizeAxis()logic and add extra shape validation in JS pad/slice utilities.
Reviewed changes
Copilot reviewed 12 out of 12 changed files in this pull request and generated 9 comments.
Show a summary per file
| File | Description |
|---|---|
| onnxruntime/core/providers/cpu/ml/linearregressor.cc | Reject mismatched coefficients size to prevent GEMM backend OOB reads. |
| onnxruntime/test/providers/cpu/ml/linearregressor_test.cc | Add regression test covering invalid coefficients size. |
| onnxruntime/core/providers/webgpu/tensor/slice.cc | Add axis range checking and output-dimension validation in WebGPU Slice. |
| onnxruntime/core/providers/webgpu/tensor/pad.cc | Add output-dimension negative check in WebGPU Pad. |
| onnxruntime/core/providers/cuda/tensor/split.cc | Reject negative values in runtime split input. |
| onnxruntime/core/providers/cuda/tensor/pad.cc | Add output-dimension negative check in CUDA Pad. |
| onnxruntime/core/providers/cpu/tensor/split.h | Reject negative values in runtime split input (CPU). |
| onnxruntime/core/providers/cpu/tensor/pad.cc | Add output-dimension negative checks in CPU Pad (reshaped + final dims). |
| js/web/lib/wasm/jsep/webgpu/ops/slice.ts | Add JS WebGPU slice output-dimension validation. |
| js/web/lib/wasm/jsep/util.ts | Fix axis normalization condition; add negative-dimension check in pad shape. |
| js/web/lib/onnxjs/util.ts | Fix axis normalization condition; add negative-dimension check in pad shape. |
| js/web/lib/onnxjs/backends/webgl/ops/slice.ts | Add JS WebGL slice output-dimension validation. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
bmehta001
force-pushed
the
fix/linearregressor-heap-oob-read
branch
from
ac9a4f8 to
913cbb9
Compare
LinearRegressor treats the coefficients attribute as a [num_targets, num_features] matrix and passes it directly to MLAS SGEMM. However, num_features is derived from the input tensor at runtime, and no validation ensured coefficients.size() == num_targets * num_features. A malformed model could provide fewer coefficients than expected, causing MlasSgemmTransposePackB to read past the buffer boundary. Add a size check after num_features is computed but before the GEMM dispatch to reject mismatched coefficients with a clear error message. Files changed: - onnxruntime/core/providers/cpu/ml/linearregressor.cc - onnxruntime/test/providers/cpu/ml/linearregressor_test.cc Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
bmehta001
force-pushed
the
fix/linearregressor-heap-oob-read
branch
from
913cbb9 to
6191529
Compare
bmehta001
changed the title
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
bmehta001
force-pushed
the
fix/linearregressor-heap-oob-read
branch
from
6191529 to
33b8291
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| size_t expected_coefficients_size = 0; | ||
| ORT_TRY { | ||
| expected_coefficients_size = SafeInt<size_t>(num_targets_) * static_cast<size_t>(num_features); | ||
| } | ||
| ORT_CATCH(...) { | ||
| return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, | ||
| "LinearRegressor: coefficients attribute size (", coefficients_.size(), | ||
| ") does not match targets (", num_targets_, | ||
| ") * input features (", num_features, ")"); | ||
| } |
There was a problem hiding this comment.
In ORT_NO_EXCEPTIONS builds, ORT_TRY/ORT_CATCH are compiled to if(true)/else if(false), and SafeInt overflow calls ORT_THROW which aborts. That means the overflow path here will terminate the process instead of returning INVALID_ARGUMENT, which is undesirable for malformed/untrusted models. Consider using a non-throwing overflow check (e.g., IAllocator::CalcMemSizeForArray(static_cast<size_t>(num_targets_), static_cast<size_t>(num_features), &expected_coefficients_size) with size=1) or an explicit num_targets_ > 0 && num_features <= max/num_targets_ style check to keep behavior consistent across exception/no-exception builds.
Address review feedback: ORT_TRY/ORT_CATCH is a no-op in ORT_NO_EXCEPTIONS builds, so SafeInt overflow would abort the process. Replace with an explicit non-throwing overflow check using std::numeric_limits that works consistently across all build configs. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
bmehta001
force-pushed
the
fix/linearregressor-heap-oob-read
branch
from
33b8291 to
494c070
Compare
Description
Fix a heap-buffer-overflow (out-of-bounds read) in
LinearRegressortriggered when thecoefficientsattribute is smaller thannum_targets * num_features.Root Cause
LinearRegressortreats thecoefficientsattribute as a[num_targets, num_features]matrix and passes it directly to MLAS SGEMM viaGemm<T>::ComputeGemm. However,num_featuresis derived from the input tensor shape at runtime, and there was no validation thatcoefficients.size() == num_targets * num_features. A malformed model could provide fewer coefficients than expected, causingMlasSgemmTransposePackBto read past the buffer boundary.Fix
Added a size check in
Compute()afternum_featuresis derived but before the GEMM dispatch:SafeIntfor overflow-safe multiplicationnum_targetsis non-negativeINVALID_ARGUMENTstatus with a clear error message on mismatchTesting
CoefficientsSizeMismatch: verifies undersized coefficients are rejected with a clear errorMotivation and Context
Security fix — prevents out-of-bounds heap reads when processing untrusted ONNX models with malformed
LinearRegressorattributes on the CPUExecutionProvider.